package com.shiyi.shiyiaiagent.app;


import com.shiyi.shiyiaiagent.advisor.MyLoggerAdvisor;
import com.shiyi.shiyiaiagent.advisor.ReReadingAdvisor;
import com.shiyi.shiyiaiagent.chatMemory.FileBasedChatMemory;
import com.shiyi.shiyiaiagent.chatMemory.MySQLChatMemory;
import com.shiyi.shiyiaiagent.rag.LoveAppCustomAdvisorFactory;
import com.shiyi.shiyiaiagent.rag.LoveAppVectorStoreConfig;
import com.shiyi.shiyiaiagent.rag.QueryReWriter;
import com.shiyi.shiyiaiagent.tools.ToolRegistration;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.QuestionAnswerAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.tool.ToolCallback;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.context.annotation.Lazy;
import org.springframework.stereotype.Component;

import java.util.List;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

@Component
@Slf4j
public class LoveApp {

  private ChatClient chatClient;

  private final ChatModel dashscopeChatModel;
  record LoveReport(String title, List<String> suggestions){};

  @Resource
  @Lazy
  private MySQLChatMemory mySQLChatMemory;

  @Resource
  private VectorStore loveAppVectorStore;

  @Resource
  private Advisor loveAppRagCloudAdvisor;

//  @Resource
//  private VectorStore pgVectorVectorStore;

  @Resource
  private QueryReWriter queryReWriter;

  @Resource
  private ToolCallback[] allTools;

  private static final String SYSTEM_PROMPT = "扮演深耕恋爱心理领域的专家。开场向用户表明身份，告知用户可倾诉恋爱难题。" +
    "围绕单身、恋爱、已婚三种状态提问：单身状态询问社交圈拓展及追求心仪对象的困扰；恋爱状态询问沟通、习惯差异引发的矛盾；" +
    "已婚状态询问家庭责任与亲属关系处理的问题。引导用户详述事情经过、对方反应及自身想法，以便给出专属解决方案。\n";

  public LoveApp(ChatModel dashscopeChatModel) {
    this.dashscopeChatModel = dashscopeChatModel;
  }

  @PostConstruct  // 将初始化逻辑移到构造后
  public void init() {
    String fileDir = System.getProperty("user.dir") + "/tmp/chat-memory";
    FileBasedChatMemory fileBasedChatMemory = new FileBasedChatMemory(fileDir);
    chatClient = ChatClient.builder(dashscopeChatModel)
      .defaultSystem(SYSTEM_PROMPT)
      .defaultAdvisors(
        new MessageChatMemoryAdvisor(fileBasedChatMemory)
      )
      .build();
  }

//  public LoveApp(ChatModel dashscopeChatModel) {
//    //初始化基于文件的对话记忆
////    String fileDir = System.getProperty("user.dir") + "/chat-memory";
////    FileBasedChatMemory fileBasedChatMemory = new FileBasedChatMemory(fileDir);
//
//    //基于mysql的对话记忆
//    chatClient = ChatClient.builder(dashscopeChatModel)
//      .defaultSystem(SYSTEM_PROMPT)
//      .defaultAdvisors(
//        new MessageChatMemoryAdvisor(mySQLChatMemory)
//      )
//      .build();
//  }

  public String doChat(String message,String chatId) {

    ChatResponse chatResponse = chatClient.prompt()
      .user(message)
      .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10)
      )
      .call()
      .chatResponse();
    String answer = chatResponse.getResult().getOutput().getText();
    log.info("answer:{}", answer);
    return answer;
  }

  public LoveReport doChatWithReport(String message,String chatId) {

    LoveReport loveReport = chatClient.prompt()
      .system(SYSTEM_PROMPT + "每次对话后都要生成恋爱结果，标题为{用户名}的恋爱报告，内容为建议列表")
      .user(message)
      .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10)
      )
      .call()
      .entity(LoveReport.class);
    log.info("loveReport:{}", loveReport);
    return loveReport;
  }

  public String doChatWithRag(String message,String chatId) {
    String rewrittenMessage = queryReWriter.rewrite(message);

    ChatResponse chatResponse = chatClient.prompt()
      .user(rewrittenMessage)
      .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10)
      )
      .advisors(new MyLoggerAdvisor())
      //基于云知识库的检索增强
//      .advisors(loveAppRagCloudAdvisor)
      //基于PgVector的检索增强
//      .advisors(new QuestionAnswerAdvisor(pgVectorVectorStore))
      .advisors(new QuestionAnswerAdvisor(loveAppVectorStore))
      //基于自定义的RAG检索增强器（文档增强器 + 上下文增强器）
//      .advisors(LoveAppCustomAdvisorFactory.createLoveAppRagCustomAdvisor(loveAppVectorStore,"单身"))
      .call()
      .chatResponse();
    String answer = chatResponse.getResult().getOutput().getText();
    log.info("answer:{}", answer);
    return answer;
  }

  public String doChatWithTools(String message,String chatId) {

    ChatResponse chatResponse = chatClient.prompt()
      .user(message)
      .advisors(spec -> spec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId)
        .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10)
      )
      .advisors(new MyLoggerAdvisor())
      .tools(allTools)
      .call()
      .chatResponse();
    String answer = chatResponse.getResult().getOutput().getText();
    log.info("answer:{}", answer);
    return answer;
  }
}

